//----------------------------------*-C++-*----------------------------------//
/*!
 * \file   PowerIter.cc
 * \author Jeremy Roberts
 * \date   11/24/2010
 * \brief  Member definitions of base class PowerIter
 * \note   Copyright (C) 2010 Jeremy Roberts
 */
//---------------------------------------------------------------------------//
// $Rev::                                               $:Rev of last commit
// $Author::                                            $:Author of last commit
// $Date::                                              $:Date of last commit
//---------------------------------------------------------------------------//

#include <iostream>
#include <cmath>
#include "LinAlg.hh"
#include "GlobalInput.hh"
#include "ResponseFunctionServer.hh"
#include "ResponseMatrix.hh"
#include "ResponseMatrixFull.hh"
#include "AbsorptionResponse.hh"
#include "FissionResponse.hh"
#include "LeakageResponse.hh"
#include "ConnectMatrix.hh"
#include "Connect2dCart.hh"
#include "GlobalProblem.hh"
#include "GlobalSolver.hh"
#include "PowerIter.hh"
#include "InvItShell.hh"

//---------------------------------------------------------------------------//
// CONSTRUCTOR
//---------------------------------------------------------------------------//
PowerIter::PowerIter( GlobalProblem &pr, GlobalInput &in )
 : GlobalSolver(pr,in)
{
    std::cout << " CONSTRUCTING PowerIter " << std::endl;
    // nothing more here for now
}

//---------------------------------------------------------------------------//
// DESTRUCTOR
//---------------------------------------------------------------------------//
PowerIter::~PowerIter()
{
    // nothing more here right now, as M is destroyed explicitly in Release()
    return; 
}

//---------------------------------------------------------------------------//
/*!
 * \brief This is the power iteration solver.
 *
 * It solves
 *   \f[
 *     \mathbf{M}\mathbf{R}(k^{(m)}) J_{-}^{(n,m)} = J_{-}^{(n+1,m)} 
 *   \f]
 * for an inner iteration \f$ n \f$ within an outer iteration \f$ m \f$. 
 * The outer iteration is defined by the eigenvalue update
 *   \f[
 *     k^{(m+1)} = \frac{ \mathbf{F}(k^{(m)})J_{-}^{(n,m)} } { \mathbf{L}(k^{(m)})J_{-}^{(n,m)} } 
 *   \f]
 * for fission operator \f$ \mathbf{F} \f$ and loss (absorption plus leakage) 
 * operator \f$ \mathbf{L} \f$, and where \f$ n \f$ is the number of inner 
 * iterations required for convergence on \f$ J_{-} \f$.
 */
void PowerIter::Solve()
{

    scalar keff = input->keff;  
    scalar lambda = 1.0;

    scalar one  = 1.0;
    scalar zero = 0.0;

    SermentVector Jtmp1( input->degfree );
    SermentVector Jtmp2( input->degfree );
    SermentVector Jtmp3( input->degfree );
    SermentVector *Jo, *J1, *J2, *J3, *Jtmp;
    Jo  = &problem->J;
    J1  = &Jtmp1;
    J2  = &Jtmp2;
    J3  = &Jtmp3;

    for ( integer i=0; i<input->degfree; i++ )
    {
        Jo->insertVal( i, zero );
    }
    // uniform zeroth order initial guess
    for ( integer i=0; i<input->degfree; 
          i = i + (input->spaceord+1)*(input->angleord+1)*input->numgroups )
    {
        Jo->insertVal( i, one );
    }    

    scalar loss=0, gain=0, absorption=0, leakage=0, keffo=0, keffoo=0;

    scalar scale; 
    scale = 1 / sqrt( Jo->vecDot(*Jo) );
    Jo->vecScale(scale); // initial guess now has unity length


    bool aitkenOn = true;          // only used to view aitken compared to non
    bool steffensenOn = false;     // actually plugs aitken value into rf's
    integer maxIt      = 500;
    integer innerCheck = maxIt/1000;
    integer it =0, itinner = 0, innertot = 0; scalar norm=0;
    integer linits = 0;

integer ittype = input->maxit-20;
if ( ittype == 0 )
{
    // - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    // POWER ITERATION
    // - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    cout << " beginning power iterations.... ";
    if (steffensenOn) cout << " with Steffensen ";
    cout << endl;
    for ( it=1; it<=input->maxit; it++)
    {

        // update all response operators -- not updated here, since they
        // are updated in normResid
        problem->R.updateData( keff );
        problem->L.updateData( keff );
        problem->A.updateData( keff );
        problem->F.updateData( keff ); 
        itinner = 0;
        for (integer i=0; i < maxIt; i++)
        {
            problem->R.matVec(*Jo,*J1);          // = R*J
            problem->M.matVec(*J1,*J2);          // = M*R*J
            lambda = sqrt( J2->vecDot(*J2) ); 
            J2->vecScale(1.0/lambda);
            itinner++; innertot++;
            Jtmp = J2; J2 = J1; J1 = Jo; Jo = Jtmp;  // switches Jo and J2
//            if ( itinner == innerCheck )       // check ||Jo-J1||, ie || Jn - Jn+1 ||
//            {
//                J1->vecAYPV(-1.0,*Jo);
//                scalar tmpdj = sqrt( J1->vecDot(*J1) );
//                if ( tmpdj <= 1e-10 ) break;
//                itinner=0;
//            }

        }

        // compute gains and losses
        gain       = problem->F.vecDot(*Jo);
        absorption = problem->A.vecDot(*Jo);
        leakage    = problem->L.computeLeakage(*Jo);
        loss       = absorption+leakage;
        keffoo     = keffo; // keep oldest value
        keffo      = keff;  // keep second oldest value
        keff       = gain/loss;
        if ( it > 2 and steffensenOn ) keff = Aitken(keffoo,keffo,keff);

        norm  = normResid(Jo,keff,lambda);

        if ( steffensenOn or !aitkenOn )
        {
        printf ("  %3i  POWER ITERATION function norm  %8.6g lambda %12.9f keff %12.9f cumulative inners %8i \n", 
                it, norm, lambda, keff, innertot );
        }
        else
        {
        printf ("  %3i  POWER ITERATION function norm  %8.6g lambda %12.9f keff %12.9f keffAit %12.9f cumulative inners %8i \n", 
                it, norm, lambda, keff, Aitken(keffoo,keffo,keff), innertot );
        }

		// test fissionrates
		//fissionRates(Jo);

		// the norm of the residual is sufficiently small
        if ( norm < 1e-8 ) break;

		// warning about convergence
        if ( it == input->maxit-1 )
            cout << "****warning, reached max outers w/o convergence" << endl;   

    }
}
else
{
    cout << " beginning inverse iterations.... " << endl;
    // - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    // INVERSE ITERATION -- could not get to work -- looks like a cond. problem
    // - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        integer nzz[ input->degfree ];
        for (integer i=0; i < input->degfree; i++)
            nzz[i] = input->degfree/input->numel + 2;
        const integer *nzzp = nzz;
    Mat MR; Mat Rcsr;
    SermentMatrixCRS pc( input->degfree, input->degfree, input->degfree/input->numel, nzzp );
    scalar tpc = MPI_Wtime();
    // MR-I*lambda approximation for an ILU
    problem->R.updateData( keff ); 
    MatConvert( problem->R.M, MATAIJ, MAT_INITIAL_MATRIX, &Rcsr );
    MatMatMult( problem->M.M, Rcsr, MAT_INITIAL_MATRIX, 1.0, &MR );
    SermentVector D1(input->degfree);
    VecSet( D1.V, -1.0 );
//    for (integer i = 0; i < input->degfree; i++)
//        D1.insertVal( i, -1.0 );
    //MatDiagonalSet( MR, D1.V, ADD_VALUES ); D1.releaseMe();
    SermentVector D2( input->degfree );
    MatGetDiagonal( MR, D2.V );
    D2.vecAYPV( 1.0, D1 );
    scalar *d_a;
    VecGetArray( D2.V, &d_a );
    integer row, ncols;
    const integer *cols;
    const scalar *vals; 
    scalar tmpval = 0.0;
    for ( row = 0; row < input->degfree; row++ )
    {
        MatGetRow( MR, row, &ncols, &cols, &vals );
        MatSetValues(pc.M,1,&row,ncols,cols,vals,INSERT_VALUES);
        tmpval = d_a[row]; 
        MatSetValues(pc.M,1,&row,1,&row,&tmpval,INSERT_VALUES); // just insert it over a possibly inserted values of MR[i,i]
        MatRestoreRow( MR, row, &ncols, &cols, &vals );
    }
    pc.checkReady();
//    PetscViewer bview;
//    PetscViewerASCIIOpen(PETSC_COMM_WORLD,"MR.bin",&bview);
//    PetscViewerSetFormat(bview,PETSC_VIEWER_ASCII_DENSE);
//    MatView(MR,bview);
    // done
    cout << " pc build time = " << MPI_Wtime()-tpc << endl;
    InvItShell  A(input->degfree,input->degfree,PETSC_NULL,problem);
    KSP invksp;
    KSPCreate( PETSC_COMM_SELF, &invksp );
   // KSPSetTolerances(invksp,PETSC_DEFAULT,PETSC_DEFAULT,PETSC_DEFAULT,100 );
    KSPSetType(invksp,KSPGMRES);
   // KSPGMRESSetRestart(invksp,10);
    KSPSetOperators( invksp, A.M, pc.M, SAME_NONZERO_PATTERN );
    //KSPSetPreconditionerSide(invksp,PC_RIGHT);
    KSPSetFromOptions(invksp);
    integer itstmp = 0;
    for ( it=1; it<=input->maxit; it++)
    {
        problem->R.updateData( keff );
        problem->L.updateData( keff );
        problem->A.updateData( keff );
        problem->F.updateData( keff ); 

        // MR-I*lambda approximation for an ILU
//        MatConvert(problem->R.M, MATAIJ, MAT_INITIAL_MATRIX, &Rcsr );
//        MatMatMult( problem->M.M, Rcsr, MAT_INITIAL_MATRIX, 1.0, &MR );
//        SermentVector D1(input->degfree);
//        for (integer i = 0; i < input->degfree; i++)
//            D1.insertVal( i, -lambda );
//        MatDiagonalSet( MR, D1.V, ADD_VALUES ); D1.releaseMe();
//        KSPSetOperators( invksp, A.M, MR, SAME_NONZERO_PATTERN );
        // done

        itinner = 0;
        for (integer i=0; i < 2; i++)
        {
            //J1->viewMe();
            A.updateEigs(keff,lambda);
            KSPSolve( invksp, Jo->V, J1->V );  // solve(A,b,x)

            lambda = sqrt(J1->vecDot(*J1));    // temp scale
            J1->vecScale(1.0/lambda);          // normalize J1
            //J1->viewMe();
            problem->R.matVec(*J1,*J2);        // = R*J1
            problem->M.matVec(*J2,*J3);        // = M*R*J1
            lambda = sqrt(J3->vecDot(*J3));    // = J1'(M*R*J1)
            J3->vecScale(1.0/lambda);
            //cout << " lambda = " << lambda << endl;
            itinner++; innertot++;
            Jtmp = J3; J3 = J2; J2 = J1; J1 = Jo; Jo = Jtmp;  // switches Jo and J2
//            if ( itinner == innerCheck )       // check ||Jo-J1||, ie || Jn - Jn+1 ||
//            {
//                J1->vecAYPV(-1.0,*Jo);
//                scalar tmpdj = sqrt( J1->vecDot(*J1) );
//                if ( tmpdj <= 1e-8 ) break;
//                itinner=0;
//            }
            KSPGetIterationNumber(invksp, &itstmp);
            linits = linits + itstmp;

        }
        // compute gains and losses
        gain       = problem->F.vecDot(*Jo);
        absorption = problem->A.vecDot(*Jo);
        leakage    = problem->L.computeLeakage(*Jo);
        loss       = absorption+leakage;
        keffoo     = keffo; // keep oldest value
        keffo      = keff;  // keep second oldest value
        keff       = gain/loss;
        if ( it > 2 and steffensenOn ) keff = Aitken(keffoo,keffo,keff);

        norm       = normResid(Jo,keff,lambda);
        printf ("  %3i  INV ITERATION function norm  %8.6g keff %12.9f keffAit %12.9f cumulative inners %8i \n", 
                it, norm, keff, Aitken(keffoo,keffo,keff), innertot );

        if ( norm < 1e-8 ) break;
        if ( it == input->maxit-1 )
            cout << "****warning, reached max outers w/o convergence" << endl;            
    }


}

    cout << " FINAL POWER ITERATION EIGENVALUES: " << endl;  
    printf (" **** FINAL KEFF   = %12.9f \n", keff );
    printf (" **** FINAL LAMBDA = %12.9f \n", lambda );
    printf (" **** OUTER ITERATIONS  = %8i \n", it );
    printf (" **** INNER ITERATIONS  = %8i \n", innertot );
    printf (" **** LINEAR ITERATIONS = %8i \n", linits );
  //  Jo->viewMe();

    


    return;
}

//---------------------------------------------------------------------------//
/*!
 * \brief This function computes the \f$ L_2 \f$ norm of the nonlinear residual.
 *
 * The nonlinear residual is defined as
 *   \f[
 *       \mathbf{f(x)} = \left [\begin{array}{c}
 *	        (\mathbf{M}\mathbf{R}(k)-\lambda \mathbf{I}) \mathbf{J_-} \\
 *	        \mathbf{F}(k)\mathbf{J_-} - (k\mathbf{L}(k)\mathbf{J_-} ) \\
 *	        \frac{1}{2} \mathbf{J^T_-} \mathbf{J_-} - \frac{1}{2}  
 *	      \end{array} 
 *       \right ]  = \mathbf{0} \, ,
 *   \f]
 * which is the same as used in the Newton-based schemes.  The \f$ L_2 \f$
 * norm is then \f$ \sqrt{ \mathbf{f(x)}^T \mathbf{f(x)} } \f$.
 *
 */
scalar PowerIter::normResid( SermentVector *J, scalar keff, scalar lambda )
{

    // update response quantities
    problem->R.updateData( keff );
    problem->L.updateData( keff );
    problem->A.updateData( keff );
    problem->F.updateData( keff ); 

    // Fj, Fk, and Flambda are the residuals
    SermentVector Fj( input->degfree );   
    SermentVector temp( input->degfree );

    problem->R.matVec(*J,temp);
    problem->M.matVec(temp,Fj);

    Fj.vecAYPV( -lambda, *J );                     // = (M*R-lambda*I)J

    scalar normFjsq = Fj.vecDot(Fj);               // = ||(M*R-lambda*I)J||^2
    scalar normFk = problem->F.vecDot(*J) -        // = (F-kL)J
                    keff*(problem->A.vecDot(*J)+problem->L.computeLeakage(*J));
    scalar normFlambda = 0.5 - 0.5*J->vecDot(*J);  
    scalar norm = sqrt( normFjsq + pow(normFk,2.0) + pow(normFlambda,2.0) );

    // debug
    //    cout << " normFj      = " << sqrt( normFjsq ) << endl;
    //    cout << " normFk      = " << normFk << endl;
    //    cout << " normFlambda = " << normFlambda << endl;

    Fj.releaseMe();
    temp.releaseMe();

    return norm;
}

//---------------------------------------------------------------------------//
/*!
 * \brief This function computes the Aitken/Steffenson estimate.
 *
 * Aitken's \f$ \Delta^2 \f$ method is an extremely easy way to
 * accelerate convergence of a series satisifying a few constraints.  Suppose we 
 * have a sequence of estimates for the eigenvalue: 
 * \f$ k^{(n-2)} \f$, \f$ k^{(n-1)} \f$, and \f$ k^{(n)} \f$, which converge to 
 * \f$ k^* \f$ as \f$ n\to \infty \f$.  We define a new sequence
 * \f$ k'^{(n)} \f$ such that
 * \f[
     k'^{(n)} = k^{(n)} - \frac{ ( k^{(n)} - k^{(n-1)} )^2 }{ k^{(n)} - 2k^{(n-1)} + k^{(n-2)} } \, .
 * \f]
 * It can be shown that this new sequence \f$ k'^{(n)} \f$  converges to 
 * \f$ k^* \f$ faster than the sequence \f$ k^{(n)} \f$ for monotonically 
 * converging \f$ |k^{(n)}| \f$.  Moreover, this new update can be used to 
 * update the response quantities; used in this way, Aitken's becomes 
 * Steffensen's method.
 *
 */
scalar PowerIter::Aitken( scalar k0, scalar k1, scalar k2 )
{
    scalar kA = k0 - pow(k1-k0,2) / (k2 - 2.0*k1 + k0);
    return kA;
}

//---------------------------------------------------------------------------//
//                 end of PowerIter.cc
//---------------------------------------------------------------------------//

